from aws_cdk import (
    Stack,
    aws_stepfunctions as sfn,
    aws_stepfunctions_tasks as sfn_tasks,
    aws_lambda as _lambda,
    aws_iam as iam,
    aws_sns as sns,
    aws_sns_subscriptions as sns_subs,
    CfnParameter,
    Duration,
    aws_apigateway as apigw,
    aws_s3 as s3,
    RemovalPolicy,
    CfnOutput,
    Duration
)
from constructs import Construct

class ApigwStepFunctionBedrockSns(Stack):

    def __init__(self, scope: Construct, construct_id: str, **kwargs) -> None:
        super().__init__(scope, construct_id, **kwargs)

        #get email address from CDK deploy parameter to send Pre-Signed URL for generated Video
        email_param = CfnParameter(self, "EmailParameter", 
            type="String",
            description="Email address for SNS subscription"
        )

        #create IAM policy to grant Bedrock model invoke permissions
        invoke_model_policy = iam.Policy(self, "InvokeModelPolicy",
            statements=[
                iam.PolicyStatement(
                    actions=["bedrock:InvokeModel","bedrock:Get*"],
                    resources=[f"arn:aws:bedrock:{self.region}::foundation-model/*", f"arn:aws:bedrock:{self.region}:{self.account}:async-invoke/*"]
                )
            ]
        )

        #create S3 bucket to store videos and generate Pre-Signed URL
        my_bucket = s3.Bucket(
            self, 
            "VideoBucket",
            versioned=True,
            removal_policy=RemovalPolicy.DESTROY,
             auto_delete_objects=True # This will delete the bucket on stack deletion
            )
        
        # create the Lambda function that invokes Bedrock Model
        bedrock_invoke_lambda_function = _lambda.Function(self, "InvokeModelFunction",
            runtime=_lambda.Runtime.PYTHON_3_14,
            handler="invoke.handler",
            code=_lambda.Code.from_asset("./functions/functioncode.zip"),
            timeout=Duration.seconds(900),
            environment={"BUCKET_NAME": my_bucket.bucket_name}
            )
        my_bucket.grant_read_write(bedrock_invoke_lambda_function)
        invoke_model_policy.attach_to_role(bedrock_invoke_lambda_function.role)

        # create the Lambda function that fetches status of Video Generation
        fetch_status_lambda_function = _lambda.Function(self, "FetchStatusFunction",
            runtime=_lambda.Runtime.PYTHON_3_14,
            handler="fetchstatus.handler",
            code=_lambda.Code.from_asset("./functions/functioncode.zip"),
            timeout=Duration.seconds(900),
            environment={"BUCKET_NAME": my_bucket.bucket_name}
            )
        my_bucket.grant_read_write(fetch_status_lambda_function)
        invoke_model_policy.attach_to_role(fetch_status_lambda_function.role)

        #create SNS Topic and Email Subscription
        topic = sns.Topic(self, "VideoGenerationTopic")
        topic.add_subscription(sns_subs.EmailSubscription(email_param.value_as_string))

        # create role for SFN and grant permissions to invoke Lambda Functions and publish to SNS
        sfn_role = iam.Role(
            self, 'StepFunctionsRole',
            assumed_by=iam.ServicePrincipal('states.amazonaws.com')
        )
        bedrock_invoke_lambda_function.grant_invoke(sfn_role)
        fetch_status_lambda_function.grant_invoke(sfn_role)
        topic.grant_publish(sfn_role)

        # Create Step Functions tasks
        invoke_task = sfn_tasks.LambdaInvoke(
            self, 'StartVideoGeneration',
            lambda_function=bedrock_invoke_lambda_function,
            output_path='$.Payload',
        )

        wait_task = sfn.Wait(
            self, "WaitForVideoGeneration",
            time=sfn.WaitTime.duration(Duration.minutes(3))
        )

        fetch_status_task = sfn_tasks.LambdaInvoke(
            self, 'FetchStatusOfVideoGeneration',
            lambda_function=fetch_status_lambda_function,
            output_path='$.Payload',
        )

        choice_state = sfn.Choice(self, "CheckStatus")
        completed_condition = sfn.Condition.string_equals("$.Status", "Completed")
        in_progress_condition = sfn.Condition.string_equals("$.Status", "InProgress")

        send_message_task = sfn_tasks.SnsPublish(
            self, 'SNSSendMessage',
            topic=topic,
            message=sfn.TaskInput.from_json_path_at('$'),
            subject="Pre-Signed URL for Video Generated by Bedrock Nova Reel Model"
        )

        # Create Step Functions state machine
        state_machine = sfn.StateMachine(
            self, "VideoGenerationStateMachine",
            definition_body=sfn.DefinitionBody.from_chainable(
                invoke_task
                .next(wait_task)
                .next(fetch_status_task)
                .next(choice_state
                    .when(completed_condition, send_message_task)
                    .when(in_progress_condition, wait_task)
                    .otherwise(send_message_task)
                )),
            role=sfn_role,
            timeout=Duration.minutes(15)
        )

        # Create an IAM role for API Gateway to start Step Functions executions
        api_role = iam.Role(
            self, 'ApiGatewayRole',
            assumed_by=iam.ServicePrincipal('apigateway.amazonaws.com')
        )

        # Grant permission to start execution of the state machine
        state_machine.grant_start_execution(api_role)

        #create api gateway
        api = apigw.RestApi(self, "VideoGenerator",)

        #create a new resource
        invoke_resource = api.root.add_resource("generateVideo")
        invoke_integration = apigw.AwsIntegration(
            service='states',
            action='StartExecution',
            integration_http_method='POST',
            options=apigw.IntegrationOptions(
                credentials_role=api_role,
                request_templates={
                    "application/json": f'''{{
                        "input": "$util.escapeJavaScript($input.json('$'))",
                        "stateMachineArn": "{state_machine.state_machine_arn}"
                    }}'''
                },
                integration_responses=[{
                    'statusCode': '200',
                    'responseTemplates': {
                        "application/json": '''{"executionArn": "$util.parseJson($input.json('$')).executionArn", "Result":"You will receive an email with pre-Signed URL in 4-5 minutes"}'''
                    }
                }]
            )
        )
        invoke_resource.add_method('POST', invoke_integration, method_responses=[{'statusCode': '200'}])

        CfnOutput(self, "S3-Generation-Bucket", value=my_bucket.bucket_name)

